#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from typing import Any, Dict, Generic, Type, TypeVar, List

from pydantic import BaseModel
from sqlalchemy import select, update, delete, text, func
from sqlalchemy.ext.asyncio import AsyncSession
from .models import Base

ModelType = TypeVar('ModelType', bound=Base)
CreateSchemaType = TypeVar('CreateSchemaType', bound=BaseModel)
UpdateSchemaType = TypeVar('UpdateSchemaType', bound=BaseModel)


class IDList(BaseModel):
    ids: List[int]


class PageInfo:
    page: int = 0
    limit: int = 10
    order: str = ''
    sort: bool = False

    def __init__(self, order: str, page: int, limit: int, asc: bool):
        self.order = order
        self.page = page
        self.limit = limit
        self.asc = asc


class CRUDBase(Generic[ModelType, CreateSchemaType, UpdateSchemaType]):
    def __init__(self, model: Type[ModelType]):
        self.model = model

    async def get(self, db: AsyncSession, *, pk: int = None) -> ModelType:
        """
        通过主键 id 获取一条数据

        :param db:
        :param pk:
        :return:
        """

        model = await db.execute(select(self.model).where(self.model.id == pk))
        return model.scalars().first()

    async def select_page(self, db: AsyncSession, page: PageInfo, where: dict):
        rs = await db.execute(select(func.count(self.model.id)).filter_by(**where))
        total = rs.scalars().first()
        rows = []
        if total > 0:
            query = (select(self.model).filter_by(**where).offset((page.page - 1) * page.limit))
            if page.order == '':
                if page.sort:
                    query = query.order_by(self.model.id)
                else:
                    query = query.order_by(self.model.id.desc())
            else:
                if page.sort:
                    query = query.order_by(text(page.order))
                else:
                    query = query.order_by(text(page.order + ' DESC'))

            query = query.limit(page.limit)

            rs = await db.execute(query)
            rows = rs.scalars().all()
        return total, rows

    async def create(self, db: AsyncSession, data: dict, last_id: bool = False) -> None:
        """
        新增一条数据

        :param last_id:
        :param data:
        :param db:
        :return:
        """
        create_data = self.model(**data)
        db.add(create_data)
        if last_id:
            await db.commit()
            res = await db.execute(text("SELECT LAST_INSERT_ID() as last_id"))
            return res.scalars().first()
        else:
            return None

    async def update(self, db: AsyncSession, data: dict) -> int:
        """
        通过主键 id 更新一条数据

        :param db:
        :param data: 
        :return:
        """
        result = await db.execute(update(self.model).where(self.model.id == data['id']).values(**data))
        return result.rowcount

    async def delete_by_id(self, db: AsyncSession, pk: int) -> int:
        """
        通过主键 id 删除一条数据

        :param db:
        :param pk:
        :return:
        """
        result = await db.execute(delete(self.model).where(self.model.id == pk))
        return result.rowcount

    async def delete_by_ids(self, db: AsyncSession, ids: List[int]) -> int:
        """
        通过主键 id 删除一条数据

        :param db:
        :param pk:
        :return:
        """
        result = await db.execute(delete(self.model).where(self.model.id.in_(ids)))
        return result.rowcount
